from __future__ import print_function
import os
import torch
import torch.backends.cudnn as cudnn
import numpy as np
import cv2
from detection.net_slim import Slim
import torch
from itertools import product as product
from math import ceil, pi
import time
import pickle

#torch.cuda.synchronize()


class PriorBox(object):
    def __init__(self, cfg, image_size=None, phase='train'):
        super(PriorBox, self).__init__()
        self.min_sizes = cfg['min_sizes']
        self.steps = cfg['steps']
        self.clip = cfg['clip']
        self.image_size = image_size
        self.feature_maps = [
            [ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]
        self.name = "s"
        # print(self.image_size)

    def forward(self):
        anchors = []
        for k, f in enumerate(self.feature_maps):
            min_sizes = self.min_sizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / self.image_size[1]
                    s_ky = min_size / self.image_size[0]
                    dense_cx = [x * self.steps[k] / self.image_size[1]
                                for x in [j + 0.5]]
                    dense_cy = [y * self.steps[k] / self.image_size[0]
                                for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]

        # back to torch land
        print('anchors length: ', len(anchors))
        output = torch.Tensor(anchors).view(-1, 4)
        if self.clip:
            output.clamp_(max=1, min=0)
        return output


def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
    return keep


def decode(loc, priors, variances):
    boxes = torch.cat((
        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
    boxes[:, :2] -= boxes[:, 2:] / 2
    boxes[:, 2:] += boxes[:, :2]
    return boxes


def decode_landm(pre, priors, variances):
    landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 2:4] *
                        variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 4:6] *
                        variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 6:8] *
                        variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 8:10] *
                        variances[0] * priors[:, 2:],
                        ), dim=1)
    return landms


torch.set_grad_enabled(False)


def check_keys(model, pretrained_state_dict):
    ckpt_keys = set(pretrained_state_dict.keys())
    model_keys = set(model.state_dict().keys())
    used_pretrained_keys = model_keys & ckpt_keys
    unused_pretrained_keys = ckpt_keys - model_keys
    missing_keys = model_keys - ckpt_keys
    print('Missing keys:{}'.format(len(missing_keys)))
    print('Unused checkpoint keys:{}'.format(len(unused_pretrained_keys)))
    print('Used keys:{}'.format(len(used_pretrained_keys)))
    assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
    return True


def remove_prefix(state_dict, prefix):
    ''' Old style model is stored with all names of parameters sharing common prefix 'module.' '''
    #print('remove prefix \'{}\''.format(prefix))
    def f(x): return x.split(prefix, 1)[-1] if x.startswith(prefix) else x
    return {f(key): value for key, value in state_dict.items()}


def load_model(model, pretrained_path, load_to_cpu):
    #print('Loading pretrained model from {}'.format(pretrained_path))
    if load_to_cpu:
        pretrained_dict = torch.load(
            pretrained_path, map_location=lambda storage, loc: storage)
    else:
        device = torch.cuda.current_device()
        pretrained_dict = torch.load(
            pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
    if "state_dict" in pretrained_dict.keys():
        pretrained_dict = remove_prefix(
            pretrained_dict['state_dict'], 'module.')
    else:
        pretrained_dict = remove_prefix(pretrained_dict, 'module.')
    #check_keys(model, pretrained_dict)
    model.load_state_dict(pretrained_dict, strict=False)
    return model


class Model:
    def __init__(self, priors) -> None:
        self.cfg = {
            'name': 'slim',
            'min_sizes': [[10, 16, 24], [32, 48], [64, 96], [128, 192, 256]],
            'steps': [8, 16, 32, 64],
            'variance': [0.1, 0.2],
            'clip': False,
            'loc_weight': 2.0,
            'gpu_train': True,
            'batch_size': 32,
            'ngpu': 1,
            'epoch': 250,
            'decay1': 190,
            'decay2': 220,
            'image_size': 300
        }
        #self.device = torch.device('cuda')
        self.device = torch.device('cpu')

        net = Slim(cfg=self.cfg, phase='test')
        self.net = load_model(net, r'src/main/python/detection/slim_Final.pth', True)
        self.net.eval()
        print('loading detection model finished!')
        cudnn.benchmark = True
        self.net = self.net.to(self.device)
        '''

        priorbox = PriorBox(self.cfg, image_size=(360, 640))
        self.priors = priorbox.forward()
        # print(len(priors))
        with open('360x640.pkl', 'wb') as f:
            pickle.dump(self.priors, f)
        '''
        '''
        with open('360x640.pkl', 'rb') as f:
            self.priors = pickle.load(f)'''
        self.priors = priors
        self.priors = self.priors.to(self.device)
        self.prior_data = self.priors.data

    def img_predeal(self, img_org, long_side=640):  # 320 *3   180*320  360*640
        target_size = long_side
        max_size = long_side
        im_shape = img_org.shape
        im_size_min = np.min(im_shape[0:2])
        im_size_max = np.max(im_shape[0:2])
        resize = float(target_size) / float(im_size_min)
        if np.round(resize * im_size_max) > max_size:
            resize = float(max_size) / float(im_size_max)
        if resize != 1:
            img = cv2.resize(img_org, None, None, fx=resize,
                             fy=resize, interpolation=cv2.INTER_LINEAR)
            img = np.float32(img)
        else:
            img = np.float32(img_org)
        im_h, im_w, _ = img.shape
        #t4 = time.time()
        #print('resize time: ', round((t4-t3)*1000, 1))
        #t5 = time.time()
        scale = torch.Tensor(
            [img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
        img -= (104, 117, 123)
        img = img.transpose(2, 0, 1)
        img = torch.from_numpy(img).unsqueeze(0)
        img = img.to(self.device)
        scale = scale.to(self.device)
        #t6 = time.time()
        #print('trans time:', round((t6-t5)*1000, 1))
        return img, scale, im_h, im_w, resize

    def pred_img(self, img_org):
        # 图像前处理 返参(处理好的图像 缩放比例[宽,高,宽,高]  新图像高  新图像宽 缩放比例)
        img, scale, im_height, im_width, resize = self.img_predeal(img_org)

        #tt = time.time()
        loc, conf, landms = self.net(img)  # forward pass
        #print(' tui li time :', round((time.time()-tt)*1000, 1))

        boxes = decode(loc.data.squeeze(
            0), self.prior_data, self.cfg['variance'])
        boxes = boxes * scale / resize
        boxes = boxes.cpu().numpy()

        scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
        landms = decode_landm(landms.data.squeeze(
            0), self.prior_data, self.cfg['variance'])
        scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                               img.shape[3], img.shape[2], img.shape[3], img.shape[2],
                               img.shape[3], img.shape[2]])
        scale1 = scale1.to(self.device)
        landms = landms * scale1 / resize
        landms = landms.cpu().numpy()

        # ignore low scores
        inds = np.where(scores > 0.85)[0]
        # print()
        #print('len inds:', len(inds))
        boxes = boxes[inds]
        landms = landms[inds]
        scores = scores[inds]

        # keep top-K before NMS
        order = scores.argsort()[::-1]
        # order = scores.argsort()[::-1][:args.top_k]
        boxes = boxes[order]
        landms = landms[order]
        scores = scores[order]

        # do NMS

        dets = np.hstack((boxes, scores[:, np.newaxis])).astype(
            np.float32, copy=False)

        #t14 = time.time()
        keep = py_cpu_nms(dets, 0.4)
        # keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
        #t15 = time.time()
        # print()
        # print()
        #print('py_cpu_nms time', round((t15-t14)*1000, 1))
        dets = dets[keep, :]
        landms = landms[keep]
        # keep top-K faster NMS
        # dets = dets[:args.keep_top_k, :]
        # landms = landms[:args.keep_top_k, :]
        dets = np.concatenate((dets, landms), axis=1)
        #t9 = time.time()
        #print('nms time: ', round((t9-t14)*1000, 1))
        '''

        # # 0,1,2,3 ranctangle location  4
        for b in dets:

            if b[4] < 0.6:
                continue
            text = "{:.4f}".format(b[4])
            b = list(map(int, b))
            print(b)
            cv2.rectangle(img_org, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
            cx = b[0]
            cy = b[1] + 12
            cv2.putText(img_org, text, (cx, cy),
                        cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

            # landms
            cv2.circle(img_org, (b[5], b[6]), 1, (0, 0, 255), 4)
            cv2.circle(img_org, (b[7], b[8]), 1, (0, 255, 255), 4)
            cv2.circle(img_org, (b[9], b[10]), 1, (255, 0, 255), 4)
            cv2.circle(img_org, (b[11], b[12]), 1, (0, 255, 0), 4)
            cv2.circle(img_org, (b[13], b[14]), 1, (255, 0, 0), 4)
        cv2.imwrite('saveimg.jpg', img_org)'''
        return dets


with open(r'src/main/python/detection/360x640.pkl', 'rb') as f:
    priors = pickle.load(f)
det_model = Model(priors)

if __name__ == '__main__':
    img = cv2.imread('img_test/2.png')
    img = cv2.resize(img, (1920, 1080))
    # print(net.pred_img(img))
    det_model.pred_img(img)
